#include <vector>
#include<iostream>
#include<stdio.h>
#include <algorithm>
#include "scatter_elements_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"
using namespace std;

void print_arr(int32_t* array, int32_t arr_len){
    for(int32_t i=0; i<arr_len; ++i)
    {
        printf("%d ", array[i]);
        
    }
    printf("\n");
}

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    TilingData tiling;
    uint32_t sizeofdatatype;
    uint32_t totalLengthAligned;

    // 1. 获取平台信息
    uint64_t ub_size;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto socVersion = ascendcPlatform.GetSocVersion();
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);

    uint32_t typeLength = 0;
    ge::TypeUtils::GetDataTypeLength(context->GetInputDesc(0)->GetDataType(), typeLength);

    // if (typeLength==2) exit(-1);
    // 2. 获取数据信息
    int32_t dim = *(context->GetAttrs()->GetInt(0));
    const char* reduction = context->GetAttrs()->GetStr(1);
    
    int32_t mode = 0;
    
    std::cout<<reduction<<std::endl;
    if (reduction[0]=='a'){mode=1;}
    else if (reduction[0]=='m'){mode=2;}


    // if (mode==0) exit(-1);

    // printf("reduction first char :%d\n", mode);

    uint32_t totalLengthInput = context->GetInputTensor(0)->GetShapeSize();
    uint32_t totalLengthUpdate = context->GetInputTensor(1)->GetShapeSize();

    auto shapeInput = context->GetInputTensor(0)->GetOriginShape();
    auto shapeUpdate = context->GetInputTensor(1)->GetOriginShape();
    const int32_t dimNum = shapeInput.GetDimNum();

    int32_t input_shape[dimNum];
    int32_t update_shape[dimNum];

    for(int32_t i = 0; i < dimNum; i++){
        input_shape[i] = shapeInput.GetDim(i);
    }

    for(int32_t i = 0; i < dimNum; i++){
        update_shape[i] = shapeUpdate.GetDim(i);
    }

    // printf("%d %d %d %d %d\n", totalLengthInput, totalLengthUpdate, dim, mode, dimNum);
    // print_arr(input_shape, dimNum);
    // print_arr(update_shape, dimNum);
    tiling.set_input_shape(input_shape);
    tiling.set_update_shape(update_shape);
    tiling.set_totalLengthInput(totalLengthInput);
    tiling.set_totalLengthUpdate(totalLengthUpdate);


    // if (totalLengthUpdate>20000) exit(-1);
    
    tiling.set_dim(dim);
    tiling.set_mode(mode);
    tiling.set_dimNum(dimNum);

    context->SetBlockDim(1);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
                        context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t* currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
}
}

namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
static graphStatus InferDataType(gert::InferDataTypeContext* context)
{
    const auto inputDataType = context->GetInputDataType(0);
    context->SetOutputDataType(0, inputDataType);
    return ge::GRAPH_SUCCESS;
}
}

namespace ops {
class ScatterElements : public OpDef {
public:
    explicit ScatterElements(const char* name) : OpDef(name)
    {
        this->Input("var")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_UINT8})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("indices")
                .ParamType(REQUIRED)
                .DataType({ge::DT_INT32,ge::DT_INT32,ge::DT_INT32,ge::DT_INT32})
                .Format({ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND,ge::FORMAT_ND});
        
        this->Input("updates")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_UINT8})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

        this->Attr("axis").AttrType(OPTIONAL).Int(int64_t(1));
        this->Attr("reduce").AttrType(OPTIONAL).String("None");
        
        
        
        // this->Output("y")
        //     .ParamType(REQUIRED)
        //     .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_INT8})
        //     .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
        //     .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});


        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);
        this->AICore()
            .SetTiling(optiling::TilingFunc)
            .AddConfig("ascend310b");
    }
};
OP_ADD(ScatterElements);
}
